import argparse

import cv2

import models  # noqa: F401,F403
from mmdet.apis import inference_detector, init_detector
from mmdet.datasets.voc import VOCDataset


def argument_parser():
    parser = argparse.ArgumentParser(description="CenterNet detect demo")
    parser.add_argument("img", help="Image file")
    parser.add_argument("config", help="Config file")
    parser.add_argument("checkpoint", help="checkpoint file")
    args = parser.parse_args()
    return args


def visualize_bbox(img, bbox, class_name, thickness=2):
    """Visualizes a single bounding box on the image"""
    BOX_COLOR = (255, 0, 0)  # Red
    TEXT_COLOR = (255, 255, 255)  # White
    if len(bbox) >= 5:
        x_min, y_min, x_max, y_max = [int(e) for e in bbox[:-1]]
        score = bbox[-1]
        str_info = F"{class_name} {score:.2f}"
    else:
        assert len(bbox) == 4
        x_min, y_min, x_max, y_max = [int(e) for e in bbox]
        str_info = F"{class_name}"
    cv2.rectangle(img, (x_min, y_min), (x_max, y_max),
                  color=BOX_COLOR,
                  thickness=thickness)

    ((text_width, text_height), _) = cv2.getTextSize(str_info,
                                                     cv2.FONT_HERSHEY_SIMPLEX,
                                                     0.35, 1)
    cv2.rectangle(img, (x_min, y_min - int(1.3 * text_height)),
                  (x_min + text_width, y_min), BOX_COLOR, -1)
    cv2.putText(
        img,
        text=str_info,
        org=(x_min, y_min - int(0.3 * text_height)),
        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
        fontScale=0.35,
        color=TEXT_COLOR,
        lineType=cv2.LINE_AA,
    )
    return img


def main():
    args = argument_parser()
    detector = init_detector(args.config, args.checkpoint)
    results = inference_detector(detector, args.img)

    image = cv2.imread(args.img)
    for cls_index, bboxes in enumerate(results):
        class_name = VOCDataset.CLASSES[cls_index]
        for bbox in bboxes:
            if len(bbox) == 5 and bbox[-1] <= 0.2:  # score less than 0.3
                continue
            image = visualize_bbox(image, bbox, class_name)
    cv2.namedWindow("show", cv2.WINDOW_NORMAL)
    cv2.imshow("show", image)
    cv2.waitKey()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
